from pymilvus import MilvusClient,DataType
from .embedding import siliconflow_embedding
# 后面我觉得可以每次请求创建一个
def init_milvus(path:str="data/milvus_demo.db"):

    client = MilvusClient(path)

    if client.has_collection(collection_name="paper_collection"):
        client.drop_collection(collection_name="paper_collection")
    # 创建schema
    schema = MilvusClient.create_schema(
        auto_id=True,
        enable_dynamic_field=True,
    )
    schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="title", datatype=DataType.VARCHAR, max_length=65535)
    schema.add_field(field_name="abstract", datatype=DataType.VARCHAR, max_length=65535)
    schema.add_field(field_name="authors", datatype=DataType.VARCHAR, max_length=65535)
    schema.add_field(field_name="date", datatype=DataType.VARCHAR, max_length=65535)
    schema.add_field(field_name="journal", datatype=DataType.VARCHAR, max_length=65535)
    schema.add_field(field_name="url", datatype=DataType.VARCHAR, max_length=65535)
    # schema.add_field(field_name="DOI", field_type=DataType.VARCHAR, max_length=1000)
    # schema.add_field(field_name="pdf_url", field_type=DataType.VARCHAR, max_length=10000)
    schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=1024)

    # 创建索引
    index_params = client.prepare_index_params()

    index_params.add_index(
        field_name="vector", 
        index_type="FLAT",
        metric_type="COSINE"
    )

    client.create_collection(
        collection_name="paper_collection",
        schema=schema,
        index_params=index_params,
        properties={
            "collection.ttl.seconds": 46400 # 以半天为单位
        }
    )
    client.close()
    
    return path


def insert_milvus(records:list, path:str="data/milvus_demo.db"):
    client = MilvusClient(path)
    client.insert(collection_name="paper_collection",data=records)
    client.close()

def search_milvus(query:str, path:str="data/milvus_demo.db",top_n:int=30):
    client = MilvusClient(path)
    query_embedding = siliconflow_embedding([query],model_name="BAAI/bge-m3")
    results = client.search(collection_name="paper_collection",data=query_embedding,limit=top_n,output_fields=["title","abstract","authors","date","journal","url"])
    client.close()
    return results
    
    

# BAAI/bge-m3